import pymysql
import re

class mysql:
    '''
        Mysql类
    '''

    #连接
    conn = None
    #游标
    cursor = None

    def __init__(self,
                    host='127.0.0.1',
                    user='root',
                    password='root',
                    db='ztb',
                    charset='utf8mb4',
                    port=3307
                 ):
        try:
            self.conn = pymysql.connect(host=host,user=user,password=password,db=db,charset=charset,port=port,cursorclass=pymysql.cursors.DictCursor)
            self.cursor = self.conn.cursor()

            if self.conn is None or self.cursor is None:
                raise Exception()
        except:
            exit('mysql初始化失败')



    def query(self,sql,type = 1,args = None):
        '''
        查询

        :param sql:
        :param type: 0:多行获取;1:单行获取
        :param args:参数
        :return: select返回查询结果,insert返回插入ID,批量插入返回成败,其余操作返回结果
        '''

        try:

            match = re.match(r'^(insert|update|alter|delete|select){1}\s+',sql,re.I|re.M)

            if match is not None:
                match = match.group().lower().strip()

                #返回值
                res = None

                if match == 'select':
                    self.cursor.execute(sql, args)
                    #获取单行
                    if type:
                        res = self.cursor.fetchall()
                    else:
                        res = self.cursor.fetchone()
                elif match == 'insert':
                    if not args:
                        raise Exception('缺少插入数据')
                    else:
                        #单条插入
                        if isinstance(args,dict):

                            sql += ' ('
                            tmp = ' values ('

                            for v in args:
                                sql += v + ','
                                tmp += '%(' + v + ')s,'

                            sql = sql.rstrip(',') + ')' + tmp.rstrip(',') + ')'

                            self.cursor.execute(sql,args)

                            res = self.conn.insert_id()

                        elif isinstance(args,list):

                            #批量写入不使用占位符

                            key = args[0].keys()
                            if not key:
                                raise Exception

                            sql += ' (' + ','.join(key) + ') values '

                            for v in args:
                                if not v or not isinstance(v,dict):
                                    raise Exception

                                sql += '('
                                for vv in v.values():
                                    sql += '"' + str(vv) + '",'

                                sql = sql.rstrip(',') + '),'

                            sql = sql.rstrip(',')

                            res = self.cursor.execute(sql)
                        else:
                            raise Exception('仅支持列表和字典形式')

                        self.conn.commit()
                        self.conn.close()
                else:
                    #删除,修改,或其他操作
                    res = self.cursor.execute(sql, args)

                    self.conn.commit()
                    self.conn.close()
                return res
        except Exception as e:
            self.conn.rollback()
            exit(e.args)



    def close(self):
        self.cursor = None
        self.conn = None
